# agent_structured_demo.py：自定义中文提示词 + JSON 结构化工具调用 + 可视化打印
# 适配思路：不用内置 Agent 执行器，手写一个最小可用的“结构化调用”循环，
# 强制模型按 JSON 规范输出工具调用，再由本地 Python 解析并调用对应函数。

from langchain.tools import tool
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv
import os
import json
import re
from typing import Dict, Any


# =============== 环境与模型 ===============
load_dotenv()

llm = ChatOpenAI(
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    model="qwen-turbo"
)


# =============== 定义工具（Tools） ===============
@tool(description="用于查询指定城市的当前天气。参数：city（必填，城市名，示例：北京、上海），返回值：温度和天气状况。用户问城市天气必须调用此工具，不可编造数据。")
def WeatherTool(city: str) -> str:
    weather_data = {
        "北京": "温度18-25℃，晴，微风",
        "上海": "温度20-28℃，多云，东南风3级",
        "广州": "温度25-32℃，阴，有小雨"
    }
    return weather_data.get(city, f"暂未获取到{city}的天气数据")


@tool(description="用于计算加减乘除运算。参数：expression（必填，数学表达式，示例：2+3*4、10-5，无空格），返回值：计算结果。仅处理数字和+、-、*、/。")
def CalculatorTool(expression: str) -> str:
    try:
        result = eval(expression)
        return f"计算结果：{expression} = {result}"
    except Exception as e:
        return f"计算失败：{str(e)}，请检查表达式（示例：2+3，不要有空格）"


@tool(description="用于查询公司内部知识库（如年假政策、产品价格）。参数：question（必填，查询问题，示例：公司年假政策、产品A价格），返回值：知识库答案。公司相关问题必须调用此工具，不可编造。")
def RAGTool(question: str) -> str:
    rag_data = {
        "公司年假政策": "员工入职满1年5天年假，满3年10天，满5年15天。",
        "产品A的价格": "产品A基础版999元/年，高级版1999元/年，企业版需定制。",
        "加班补贴标准": "工作日20元/小时，周末30元/小时，法定节假日50元/小时。"
    }
    for key, value in rag_data.items():
        if key in question:
            return value
    return f"知识库中未找到「{question}」的信息，请联系相关部门。"


# 将工具注册到一个映射中，便于手动分发调用
REGISTERED_TOOLS: Dict[str, Any] = {
    WeatherTool.name: WeatherTool,
    CalculatorTool.name: CalculatorTool,
    RAGTool.name: RAGTool,
}


# =============== 自定义中文提示词（要求 JSON 工具调用） ===============
SYSTEM_PROMPT = (
    "你是一个严格遵守规则的中文智能Agent。你可以使用以下工具来完成任务：\n"
    "{tools}\n\n"
    "必须遵守以下调用规则：\n"
    "1) 若无需工具即可回答，请直接输出最终答案，且不要输出任何工具调用。\n"
    "2) 若需要调用工具，请严格输出一个 JSON 代码块，格式如下：\n"
    "```json\n"
    "{{\"name\": \"工具名\", \"parameters\": {{\"参数名\": \"参数值\", ...}}}}\n"
    "```\n"
    "3) JSON 代码块之外不要输出任何多余文字。\n"
    "4) 工具返回结果后，如无需继续调用其他工具，请用简洁中文给出最终答案。\n"
)


def render_tools_doc() -> str:
    lines = []
    for tool_obj in REGISTERED_TOOLS.values():
        lines.append(f"- {tool_obj.name}{tool_obj.args} - {tool_obj.description}")
    return "\n".join(lines)


def build_prompt(user_input: str) -> str:
    tools_doc = render_tools_doc()
    prompt = (
        SYSTEM_PROMPT.format(tools=tools_doc)
        + "\n用户问题：" + user_input + "\n"
        + "请按规则作答。"
    )
    return prompt


# =============== 工具调用 JSON 提取与执行 ===============
JSON_BLOCK_RE = re.compile(r"```json\s*(\{[\s\S]*?\})\s*```", re.IGNORECASE)


def try_extract_json_block(text: str) -> Dict[str, Any]:
    match = JSON_BLOCK_RE.search(text)
    if not match:
        # 兼容无代码块包裹的纯 JSON
        text = text.strip()
        if text.startswith("{") and text.endswith("}"):
            try:
                return json.loads(text)
            except Exception:
                return {}
        return {}
    try:
        return json.loads(match.group(1))
    except Exception:
        return {}


def call_tool_by_json(tool_call: Dict[str, Any]) -> str:
    name = tool_call.get("name")
    params = tool_call.get("parameters", {})
    tool_fn = REGISTERED_TOOLS.get(name)
    if tool_fn is None:
        return f"未找到工具：{name}"
    try:
        # 打印将要执行的 JSON（结构化参数）
        print("\n====== Tool JSON 调用 ======")
        print(json.dumps({"name": name, "parameters": params}, ensure_ascii=False, indent=2))
        return tool_fn.invoke(params) if hasattr(tool_fn, "invoke") else tool_fn.run(params)
    except TypeError:
        # 对于 @tool 装饰的函数，通常可以直接以关键字参数调用
        return tool_fn(**params)


# =============== 单轮最小流程（先看是否需要工具；如需要就解析 JSON 并执行） ===============
def run_once(user_input: str) -> str:
    prompt = build_prompt(user_input)
    print("\n====== 发送给模型的完整提示词 ======")
    print(prompt)

    response = llm.invoke(prompt)
    content = response.content if hasattr(response, "content") else str(response)

    # 先尝试解析工具调用 JSON
    tool_call = try_extract_json_block(content)
    if tool_call:
        # 有工具调用，执行工具
        tool_result = call_tool_by_json(tool_call)
        print("\n====== 工具返回结果 ======")
        print(tool_result)

        # 将工具结果反馈给模型，让其给出最终中文答案
        follow_up = (
            "以下是工具的返回结果，请用简洁中文给出最终答案，不要再输出 JSON：\n"
            f"{tool_result}"
        )
        final = llm.invoke(follow_up)
        final_text = final.content if hasattr(final, "content") else str(final)
        return final_text.strip()
    else:
        # 无工具调用，视为直接回答
        return content.strip()


if __name__ == "__main__":
    print("场景1：无需工具 → 什么是人工智能？")
    out1 = run_once("什么是人工智能？用简单的话解释。")
    print("答案：", out1)

    print("\n" + "=" * 50)
    print("场景2：需要工具 → 查询北京天气")
    out2 = run_once("查询北京今天的天气。")
    print("答案：", out2)

    print("\n" + "=" * 50)
    print("场景3：计算器 → 计算表达式 12/3+4*2")
    out3 = run_once("请计算 12/3+4*2 的值。")
    print("答案：", out3)


